#!/usr/bin/env python
import sys, math, random
import src.crop as cr
from mpi4py import MPI
import numpy as np
from operator import itemgetter
from util import Timer

comm = MPI.COMM_WORLD
p = comm.size
rank = comm.rank
root = 0

step = 1
done = False

eps = 0.005
crop = cr.CRoP(eps, rank, p)
k = int(1/eps)

#===============================================================================
# A = B = np.matrix([[0.0, 0.25, 0.01],
#		 [0.25, 0.25, 0.01],
#		 [0.0, 0.01, 0.0]])
#===============================================================================

A = B = np.load('../data/np_mat__10000_0.001.npy',mmap_mode='r')

print "processor %s starting" % rank
if rank == 0 :
	print "nnz ", np.count_nonzero(A)

n = 0
with Timer() as t :
	while n < len(A[0]) :
	    crop.crop(A.T[n,:], B[n,:])
	    n += 1

print 'processor  %s finished crop spacesaving %.03f sec' % (rank,t.interval)

result = crop.get_summary(type='dict')

#data = comm.gather(summary, root)

#if rank == root :
#	for i in range(p) :
#		print data[i]
#else :
#	print "Done at rank %s" % rank
#	assert data is None

#comm.Barrier()

def merge_parallel(D1, D2):
	print "merging in parallel"
	if len(D1) > len(D2) : # swap to get min length of merge loop
		temp = D1
		D1, D2 = D2, temp

	for key,item in D1.iteritems() :
		if D2.has_key(key) :
			item = {key:(key, D2[key][1] + item[1], D2[jey][2]+item[2])}
		else :
			item = {key:item}
		D2.update(item)

	if len(D2) > k :
		D = sorted(D2.values(), key=itemgetter(1), reverse=True)
		kth = D[k][1]
		D2.clear()
		for item in D :
			if item[1] >= kth :
				D2[item[0]] = (item[0],item[1],item[2])
	return D2

while not done :
	if rank % 2 ** step != 0:
		send_msg = {"summary": result, "step": step}
		_dest = rank - 2 ** (step - 1)
		data = comm.issend(send_msg, dest=_dest)
		print rank, "sending to ", _dest
		done = True
	else :
		_source = rank + 2 ** (step - 1)
		recv_msg = comm.recv(source=_source)
#		print rank, "recv from %s at step %s: %s" % (_source, str(step), recv_msg["msg"])
		step = int(recv_msg["step"]) + 1
		result = merge_parallel(result, recv_msg["summary"])
		if step <= math.log(p, 2):
			if rank != root :
				_dest = rank - 2 ** (step - 1)
				send_msg = {"summary": result, "step": step}
				print rank, "sending to ", _dest
				data = comm.send(send_msg, dest=_dest)
		else :
			done = True
if rank == 0 :
#	print "Exact\n ", np.dot(A,B)
#	print "Rank %s. \nData: %s" % (rank, result)
#	C = np.ndarray(shape=(len(A),len(A)))
#	for key,el in result.iteritems() :
#		C[key[0]][key[1]] = el[1]
#	print C
	print "FINISHED"